#include <iostream>
#include <vector>
#include <random>
#include <chrono>
#include "add_custom_tiling.h"
#include "add_custom.cpp"

// 模拟运行时环境
namespace runtime {
    // 模拟输入输出张量
    struct Tensor {
        std::vector<float> data;
        gert::Shape shape;
        ge::DataType dtype;
        ge::Format format;

        Tensor(const gert::Shape& s, ge::DataType dt, ge::Format fmt) : shape(s), dtype(dt), format(fmt) {
            size_t size = shape.GetShapeSize();
            data.resize(size);
        }
    };

    // 模拟运行算子
    void run_op(const Tensor& input1, const Tensor& input2, Tensor& output) {
        size_t size = input1.data.size();
        for (size_t i = 0; i < size; ++i) {
            output.data[i] = input1.data[i] + input2.data[i];
        }
    }
}

int main() {
    // 定义输入张量的形状
    gert::Shape input_shape({2, 2});

    // 创建输入张量
    runtime::Tensor input1(input_shape, ge::DT_FLOAT16, ge::FORMAT_ND);
    runtime::Tensor input2(input_shape, ge::DT_FLOAT16, ge::FORMAT_ND);

    // 使用当前时间作为随机数种子，让随机数更随机
    unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
    std::default_random_engine generator(seed);

    // 扩大随机数范围，从 -10.0 到 10.0
    std::uniform_real_distribution<float> distribution(-10.0, 10.0);

    // 随机填充输入张量的数据
    for (size_t i = 0; i < input1.data.size(); ++i) {
        input1.data[i] = distribution(generator);
        input2.data[i] = distribution(generator);
    }

    // 创建输出张量
    runtime::Tensor output(input_shape, ge::DT_FLOAT16, ge::FORMAT_ND);

    // 运行算子
    runtime::run_op(input1, input2, output);

    // 输出结果
    std::cout << "Input 1:" << std::endl;
    for (float val : input1.data) {
        std::cout << val << " ";
    }
    std::cout << std::endl;

    std::cout << "Input 2:" << std::endl;
    for (float val : input2.data) {
        std::cout << val << " ";
    }
    std::cout << std::endl;

    std::cout << "Output:" << std::endl;
    for (float val : output.data) {
        std::cout << val << " ";
    }
    std::cout << std::endl;

    return 0;
}